"""
    色调曲线矫正

    author: wxz
    date: 2021-12-3
    github: https://github.com/xinzwang
"""
import sys

import numpy as np

COLOR_MATRIX = np.asarray([[.9020, -.2890, -.0715],
                           [-.4535, 1.2436, .2348],
                           [-.0934, .1919, .7086]])


def color_correction(img, color_matrix=None):
    if color_matrix is None:
        color_matrix = COLOR_MATRIX

    m = cam2rgb(color_matrix)

    out = np.zeros_like(img)
    out[:, :, 0] = img[:, :, 0] * m[0, 0] + img[:, :, 1] * m[0, 1] + img[:, :, 2] * m[0, 2]
    out[:, :, 1] = img[:, :, 0] * m[1, 0] + img[:, :, 1] * m[1, 1] + img[:, :, 2] * m[1, 2]
    out[:, :, 2] = img[:, :, 0] * m[2, 0] + img[:, :, 1] * m[2, 1] + img[:, :, 2] * m[2, 2]

    out = np.clip(out, 0.0, 1.0)

    return out


def cam2rgb(color_matrix):
    # matric multiplication
    rgb2cam = np.dot(color_matrix, rgb2xyz())

    # make sum of each row to be 1.0, necessary to preserve white balance
    # basically divice each value by its row wise sum
    rgb2cam = np.divide(rgb2cam, np.reshape(np.sum(rgb2cam, 1), [3, 1]))

    # - inverse the matrix to get cam2rgb.
    # - cam2rgb should also have the characteristic that sum of each row
    # equal to 1.0 to preserve white balance
    # - check if rgb2cam is invertible by checking the condition of
    # rgb2cam. If rgb2cam is singular it will give a warning and
    # return an identiry matrix
    if np.linalg.cond(rgb2cam) < (1 / sys.float_info.epsilon):
        return np.linalg.inv(rgb2cam)  # this is cam2rgb / color correction matrix
    else:
        print("Warning! matrix not invertible.")
        return np.identity(3, dtype=np.float32)


def rgb2xyz(color_space='srgb', illuminant='d65'):
    # Objective: get the rgb2xyz matrix dependin on the output color space
    #            and the illuminant
    # Source: http://www.brucelindbloom.com/index.html?Eqn_RGB_XYZ_Matrix.html

    matrix = {
        'srgb': {
            'd65': [[.4124564, .3575761, .1804375],
                    [.2126729, .7151522, .0721750],
                    [.0193339, .1191920, .9503041]],
            'd60': [[.4360747, .3850649, .1430804],
                    [.2225045, .7168786, .0606169],
                    [.0139322, .0971045, .7141733]]
        },
        'adobe-rgb-1998': {
            'd65': [[.5767309, .1855540, .1881852],
                    [.2973769, .6273491, .0752741],
                    [.0270343, .0706872, .9911085]],
            'd50': [[.6097559, .2052401, .1492240],
                    [.3111242, .6256560, .0632197],
                    [.0194811, .0608902, .7448387]]
        }
    }

    if color_space in matrix.keys():
        if illuminant in matrix[color_space].keys():
            return matrix[color_space][illuminant]
        else:
            print("for now, illuminant must be d65 or d50")
            return
    else:
        print("for now, color_space must be srgb or adobe-rgb-1998")
        return
